'''
无向，有向。 有无负权都可以。可以有环
但不能有负环(a->b->c->a)一个环上的权值+起来 < 0
'''
import heapq

import numpy as np


def floyd(graph, start):
    # demo为有向(无向可以在函数外设置)有权图
    n = len(graph)
    distances = np.full(shape=(n, n), fill_value=np.inf)
    np.fill_diagonal(distances, 0)
    for u in range(n):
        for v, w in graph[u]:
            distances[u][v] = w
    for m in range(n):
        for u in range(n):
            for v in range(n):
                if distances[u][m] != np.inf and distances[m][v] != np.inf \
                        and distances[u][m] + distances[m][v] < distances[u][v]:
                    distances[u][v] = distances[u][m] + distances[m][v]
    return distances[start]


# 不能有负权
def compare_dijkstra(graph, start):
    n = len(graph)
    distance = [float('inf')] * len(graph)
    distance[start] = 0
    visited = [False] * n

    heap = [(0, start)]
    while heap:
        cost, u = heapq.heappop(heap)
        if not visited[u]:
            visited[u] = True
            for v, w in graph[u]:
                if not visited[v] and cost + w < distance[v]:
                    distance[v] = cost + w
                    heapq.heappush(heap, (distance[v], v))
    return distance


if __name__ == '__main__':
    graph = [
        [(1, 1), (2, 4)],
        [(2, 2), (3, 5)],
        [(3, 1)],
        []
    ]
    print(floyd(graph, 0))
    print(compare_dijkstra(graph, 0))
